#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright (C) 2024 reinovo, Inc. All Rights Reserved 
#
# @Time    : 2024/03/14 上午10:24
# @Author  : hmm
# @Email   : liuyuhang0531@foxmail.com
# @File    : handcan_calibration.py
# @Software: PyCharm


import cv2 as cv
import numpy as np
import apriltag
from math import *
from uarm.wrapper import SwiftAPI
import matplotlib.pyplot as plt
import argparse

parse = argparse.ArgumentParser(description="手眼标定程序")
parse.add_argument('--source', help = 'Input Image Source')
parse.add_argument('--camera_param',default = 0, help = 'Camera internal parameter matrix file')
parse.add_argument('--obj_pos', type=str,default='None',help='A file that records the 3D position of the tags.')
parse.add_argument('--save_path', type=str,default='../data/eyehand_Matrix.npz',help='A file that records the 3D position of the tags.')
# parse.add_argument('--2d_position_npflie', default='../data/base2mark_pose.npz', type=str, help='A file that records the 2D position of the tags.')

args = parse.parse_args()

swift = SwiftAPI()
swift.reset()

mtx = np.load(args.camera_param)['mtx']
dist = np.load(args.camera_param)['dist']
source = args.source


if args.obj_pos == "None":
    base2tag_p = np.array([[200., 15., 0.],
                           [200., -15., 0.],
                           [170., 15, 0.],
                           [170., -15, 0.]])
else:
    base2tag_p = np.load(args.obj_pos)['pose']


    
# 定义一个检测器（使用字典“tag36h10”）
at_detector = apriltag.Detector(apriltag.DetectorOptions(families='tag36h10'))

def eulerAngles2RotationMatrix(theta):
    """
    欧拉角转换旋转矩阵
    :param theta: 欧拉角角度
    :return: 旋转矩阵
    """
    R_x = np.array([[1, 0, 0],
                    [0, cos(theta[0]), sin(theta[0])],
                    [0, -sin(theta[0]), cos(theta[0])]
                    ])
    R_y = np.array([[cos(theta[1]), 0, sin(theta[1])],
                    [0, 1, 0],
                    [-sin(theta[1]), 0, cos(theta[1])]
                    ])
    R_z = np.array([[cos(theta[2]), -sin(theta[2]),0],
                    [sin(theta[2]), cos(theta[2]),0],
                    [0, 0, 1]
                    ])
    R = R_z@R_y@R_x
    return R


def rotationMatrixToEulerAngles(R):
    """
    旋转矩阵转换欧拉角
    :param R: 旋转矩阵
    :return: 欧拉角
    """
    sy = sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])
    singular = sy < 1e-6
    if not singular:
        x = atan2(R[2, 1], R[2, 2])
        y = atan2(-R[2, 0], sy)
        z = atan2(R[1, 0], R[0, 0])
    else:
        x = atan2(-R[1, 2], R[1, 1])
        y = atan2(-R[2, 0], sy)
        z = 0
    return np.array([x, y, z])

def get_tag2end(point):
    """
    获取二维码到手臂末端
    :param point:
    :return:
    """

    # 计算绕z轴旋转角
    obj= point[:,:2]
    delta = (obj[2]-obj[3])
    r = sqrt(delta@delta.T)
    angle = np.arcsin(delta[0]/r)
    # 计算变换后的坐标
    position = swift.get_position()
    R=eulerAngles2RotationMatrix([0,0,angle])
    t = R @ np.array(position) - R @ point[3]
    return R,t


def get_img2tag(image,tag_size):
    img2tag_p = []
    cam_params = [mtx[0, 0], mtx[1, 1], mtx[0, 2], mtx[1, 2]]
    gray = cv.cvtColor(image,cv.COLOR_BGR2GRAY)
    tags = at_detector.detect(gray)
    for tag in tags:
        M, e1, e2 = at_detector.detection_pose(tag, cam_params, tag_size)
        img2tag_p.append([tag.tag_id,tag.center,M])
        print(tag.center)
    return img2tag_p


def get_tag2cam(tags):
    ret = False
    t = []
    R = []
    obj_point = np.array([[30, 30, 0],
                        [30, 0, 0],
                        [0, 30, 0],
                        [0, 0, 0]], dtype=float)
    obj_id = [0,1,2,3]
    pix_point = np.zeros((4, 2))
    for tag in tags:
        for id in obj_id:
            if tag[0] == id:
                pix_point[id, :] = tag[1]
                obj_id.remove(id)
    if len(obj_id)==0:
        pix_point_2 = np.array(pix_point,dtype=np.float32)
        obj_point_2 = np.array(obj_point,dtype=np.float32)
        print(type(pix_point_2))
        ret, rvecs, tvecs = cv.solvePnP(obj_point_2, pix_point_2, mtx, dist)
        R = cv.Rodrigues(rvecs)[0]
        t = tvecs.T
        ret = True
    return ret,R,t

def get_base2end():
    angle = swift.get_polar()[1]
    position = swift.get_position()
    R = eulerAngles2RotationMatrix([0,0,(angle-90)*pi/180])
    t = position
    return R,t

def get_base2tag_p(tags,id,R,t):
    ret = False
    dst = []
    for tag in tags:
        if tag[0]==id:
            pix2img_t = np.linalg.inv(mtx) @ np.insert(tag[1], 2, 1)
            img2tag_t = (tag[2] @ np.insert(pix2img_t, 3, 1))[:3]
            base2end_R, base2end_t = get_base2end()
            dst = base2end_t + base2end_R @ R @ img2tag_t + base2end_R @ -R@t
            ret = True
            print(dst)
    return ret,dst

if __name__ == "__main__":
    cap = cv.VideoCapture(4)
    plt.ion()
    fig = plt.figure()
    print(base2tag_p)
    while cap.isOpened():
        ret,frame = cap.read()
        cv.imshow("frame",frame)
        key = cv.waitKey(10)
        if key == 27:
            break
        elif key == ord("s"):
            tags = get_img2tag(frame,20)
            ret, tag2cam_R, tag2cam_t = get_tag2cam(tags)
            tag2end_R, tag2end_t = get_tag2end(base2tag_p)
            if ret is True:
                base2end_R, base2end_t = get_base2end()
                cam2end_R = tag2cam_R.T@tag2end_R@base2end_R
                cam2end_t = (tag2cam_R @ tag2end_t + tag2cam_t)[0]
                np.savez(args.save_path,R=cam2end_R,t=cam2end_t)
                print(f"save as {args.save_path}")
                print(f"旋转矩阵转换欧拉角{rotationMatrixToEulerAngles(cam2end_R)}")
                print(f"平移矩阵{cam2end_t}")
            else:
                print("未检测到二维码")
        elif key == ord("t"):
            print("----------------------------------------")
            with np.load(args.save_path) as X:
                R, t = [X[i] for i in ('R', 't')]
            tags = get_img2tag(frame, 20)
            dst=[]
            ax = plt.axes(projection='3d')
            ax.scatter3D(base2tag_p[:,0],base2tag_p[:,1],base2tag_p[:,2],c='r')
            for i in range(4):
                base2tag_p[i]
                ret,dst_p = get_base2tag_p(tags, i, R, t)
                if ret is True:
                    ax.scatter3D(dst_p[0],dst_p[1],dst_p[2],c='b')
            plt.pause(0.1)
        elif key == ord("d"):
            swift.set_servo_detach()
        elif key == ord("a"):
            swift.set_servo_attach()
    plt.close(fig)
    cap.release()
    cv.destroyAllWindows()
            
